import os
import timm
import torch
import numpy as np
from torchvision import transforms, datasets
from torch.utils.data import DataLoader, random_split
from torch.utils.tensorboard import SummaryWriter
from tqdm import tqdm
import torch.nn as nn
import torch.optim as optim
from torch.optim import lr_scheduler

# 设置HuggingFace镜像源
os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'

# 设备配置
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print(f'Using device: {device}')

# 数据增强配置
train_transform = transforms.Compose([
    transforms.RandomResizedCrop(224),
    transforms.RandomHorizontalFlip(),
    transforms.RandomRotation(15),
    transforms.ColorJitter(0.2, 0.2, 0.2),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

val_transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

def main():
    # 数据集路径
    dataset_dir = 'train'
    assert os.path.exists(dataset_dir), "数据集路径不存在！"

    # 加载完整数据集
    full_dataset = datasets.ImageFolder(dataset_dir)

    # 划分数据集
    train_ratio = 0.8
    val_ratio = 0.2
    total_size = len(full_dataset)
    train_size = int(train_ratio * total_size)
    val_size = total_size - train_size

    train_dataset, val_dataset = random_split(
        full_dataset,
        [train_size, val_size],
        generator=torch.Generator().manual_seed(42)
    )

    # 应用不同的transform
    train_dataset.dataset.transform = train_transform
    val_dataset.dataset.transform = val_transform

    # 类别信息
    class_names = full_dataset.classes
    n_class = len(class_names)
    print(f'发现 {n_class} 个类别')

    # 保存类别映射
    idx_to_labels = {i: name for i, name in enumerate(class_names)}
    np.save('idx_to_labels.npy', idx_to_labels)

    # 数据加载参数（Windows建议设置num_workers=0）
    BATCH_SIZE = 64
    num_workers = 4 if os.name == 'posix' else 0  # Windows设为0
    
    train_loader = DataLoader(
        train_dataset,
        batch_size=BATCH_SIZE,
        shuffle=True,
        num_workers=num_workers,
        pin_memory=True,
        persistent_workers=num_workers > 0  # 动态设置
    )

    val_loader = DataLoader(
        val_dataset,
        batch_size=BATCH_SIZE,
        shuffle=False,
        num_workers=num_workers,
        pin_memory=True,
        persistent_workers=num_workers > 0  # 动态设置
    )

    # 初始化模型
    model = timm.create_model('resnest50d', pretrained=True)
    num_features = model.fc.in_features
    model.fc = nn.Linear(num_features, n_class)
    model = model.to(device)

    # 优化器配置
    optimizer = optim.AdamW(model.parameters(), lr=1e-4)
    scheduler = lr_scheduler.CosineAnnealingLR(optimizer, T_max=20)
    criterion = nn.CrossEntropyLoss()

    # 训练参数
    EPOCHS = 50
    best_acc = 0.0
    writer = SummaryWriter('logs_resnest')

    # 训练循环
    for epoch in range(EPOCHS):
        # 训练阶段
        model.train()
        running_loss = 0.0
        progress_bar = tqdm(train_loader, desc=f'Train Epoch {epoch+1}/{EPOCHS}')
        
        for images, labels in progress_bar:
            images = images.to(device, non_blocking=True)
            labels = labels.to(device, non_blocking=True)
            
            optimizer.zero_grad()
            
            # 混合精度训练
            with torch.cuda.amp.autocast():
                outputs = model(images)
                loss = criterion(outputs, labels)
            
            loss.backward()
            optimizer.step()
            
            running_loss += loss.item() * images.size(0)
            progress_bar.set_postfix(loss=loss.item())
        
        scheduler.step()
        epoch_loss = running_loss / len(train_dataset)
        writer.add_scalar('Loss/train', epoch_loss, epoch)

        # 验证阶段
        model.eval()
        correct = 0
        total = 0
        val_loss = 0.0
        
        with torch.no_grad():
            for images, labels in tqdm(val_loader, desc='Validating'):
                images = images.to(device, non_blocking=True)
                labels = labels.to(device, non_blocking=True)
                
                outputs = model(images)
                loss = criterion(outputs, labels)
                
                val_loss += loss.item() * images.size(0)
                _, preds = torch.max(outputs, 1)
                total += labels.size(0)
                correct += (preds == labels).sum().item()
        
        epoch_acc = correct / total
        val_loss = val_loss / len(val_dataset)
        
        writer.add_scalar('Loss/val', val_loss, epoch)
        writer.add_scalar('Accuracy/val', epoch_acc, epoch)
        
        # 保存最佳模型
        if epoch_acc > best_acc:
            best_acc = epoch_acc
            torch.save({
                'epoch': epoch,
                'model_state_dict': model.state_dict(),
                'optimizer_state_dict': optimizer.state_dict(),
                'best_acc': best_acc,
            }, 'best50_model.pth')
            print(f'New best model saved with accuracy: {best_acc:.4f}')
        
        print(f'Epoch {epoch+1}/{EPOCHS} | '
              f'Train Loss: {epoch_loss:.4f} | '
              f'Val Loss: {val_loss:.4f} | '
              f'Val Acc: {epoch_acc:.4f}')

    writer.close()

if __name__ == '__main__':
    # Windows多进程必要设置
    torch.multiprocessing.freeze_support()
    
    # 设置共享内存策略
    torch.multiprocessing.set_sharing_strategy('file_system')
    
    # 减少CUDA内存碎片
    torch.backends.cudnn.benchmark = True
    
    main()